import os, sys
import pickle
import copy
from typing import List
import torch
import numpy as np
from collections import OrderedDict
from pathlib import Path
sys.path.append(os.path.join(str(Path(__file__).parent.parent), 'policy'))
from infer_agent import Agent

INIT_SCORE = 100

class PolicyManager:
    def __init__(self, model_pool_dir:str, agent:Agent, max_epoch:int=-1, learner_n:int=2) -> None:
        '''
            learner_n: oppo模型间隔
        '''
        self.model_pool = []
        self.model_path = OrderedDict()
        self.model_strength = OrderedDict()
        self.model_pool_dir = model_pool_dir
        self._model_pool_info_path = os.path.join(model_pool_dir, '_record.pkl')
        self._agent = agent
        if os.path.exists(self._model_pool_info_path):
            with open(self._model_pool_info_path, 'rb') as f:
                self.model_strength = pickle.load(f)
        self.agent_dict = {}
        self.max_epoch = max_epoch
        self.learner_n = learner_n
        for i in range(4):
            self.agent_dict[i] = self._agent[i]

    def search_model(self):
        if not os.path.exists(self.model_pool_dir):
            return
        # 初始化带有模型
        for _file in os.listdir(self.model_pool_dir):
            if _file.endswith('pth'):
                if _file not in self.model_pool:
                    self.model_pool.append(_file)
                    self.model_path[_file] = os.path.join(self.model_pool_dir, _file)
                    # 初始分
                    self.model_strength[_file] = INIT_SCORE

    def init_agents(self, oppo_idxs: List) -> List[Agent]:
        state_dict_list = self.chose_agent()
        for _idx in oppo_idxs:
            if len(state_dict_list) > 0:
                agent_model = state_dict_list.pop()
                self.agent_dict[_idx].reset(agent_model)
        return self.agent_dict

    def push_outcome(self, info:dict):
        ...

    def record_outcomes(self):
        '''
            记录对战结果
        '''
        with open(self._model_pool_info_path, 'wb') as f:
            pickle.dump(self.model_strength, f)

    def chose_agent(self) -> str:
        '''
            根据某些策略来返回一个智能体路径
        '''
        self.search_model()
        state_dict_list = []
        if len(self.model_pool) >= 3* self.learner_n:
            if self.max_epoch != -1:
                model_epoch_dict = {m: float(m.split('_')[1]) for m in self.model_pool if float(m.split('_')[1])<=self.max_epoch and float(m.split('_')[1]) % self.learner_n==0}
            else:
                model_epoch_dict = {m: float(m.split('_')[1]) for m in self.model_pool if float(m.split('_')[1]) % self.learner_n==0}
            _chose_probs = np.array(list(model_epoch_dict.values()))
            _chose_probs = _chose_probs/sum(_chose_probs)
            _model_names = np.random.choice(list(model_epoch_dict.keys()), p=_chose_probs,size=3,replace=False).tolist()
            state_dict_list = [torch.load(self.model_path[_model]) for _model in _model_names]
        return state_dict_list


if __name__ == '__main__':
    policy_manager = PolicyManager('/home/dq/jd/Competition_AAMAS2023/algo/logs/2024_03_13-14_44_33', None, 3, 40)
    policy_manager.chose_agent(0)
